# -*- coding: utf-8 -*-
"""
ERA5数据解析器

此模块提供用于解析ERA5 NetCDF格式气象数据的类。
采用面向对象设计，支持不同类型的ERA5数据格式（日降水量和小时降水量）。
"""

import datetime
import os
import re
import xarray as xr
import pandas as pd
from abc import ABC, abstractmethod
from typing import Optional, Tuple, List


class ERA5Parser(ABC):
    """
    ERA5数据解析器的抽象基类
    定义了所有ERA5解析器必须实现的接口
    """

    def __init__(self, dataset: xr.Dataset, file_path: Optional[str] = None):
        """
        初始化解析器

        Args:
            dataset: 包含ERA5数据的xarray数据集
            file_path: 数据文件路径，用于提取年份和月份信息
        """
        self.dataset = dataset
        self.time_dim = self.get_time_dimension()
        self.file_path = file_path
        self.year = None
        self.month = None
        
        if file_path:
            self._parse_file_info(file_path)

    @abstractmethod
    def get_time_dimension(self) -> str:
        """
        获取时间维度名称

        Returns:
            时间维度的名称
        """
        pass

    def _parse_file_info(self, file_path: str) -> None:
        """
        从文件路径中解析年份和月份信息
        
        Args:
            file_path: 数据文件路径
        """
        filename = os.path.basename(file_path)

        # 尝试匹配日降水量文件格式 (ERA5.P.daily.1960.nc)
        daily_match = re.search(r'ERA5\.P\.daily\.(\d{4})\.nc', filename)
        if daily_match:
            self.year = int(daily_match.group(1))
            return

        # 尝试匹配小时降水量文件格式 (ERA5-P.hourly.1940.09.nc)
        hourly_match = re.search(r'ERA5-P\.hourly\.(\d{4})\.(\d{2})\.nc', filename)
        if hourly_match:
            self.year = int(hourly_match.group(1))
            self.month = int(hourly_match.group(2))
            return

        # 尝试更通用的模式
        year_match = re.search(r'\.(\d{4})\.|\.(\d{4})$', filename)
        if year_match:
            self.year = int(year_match.group(1) or year_match.group(2))

    def format_date(self, time_value) -> str:
        """
        根据解析器类型格式化日期
        
        Args:
            time_value: 时间值
            
        Returns:
            格式化的日期字符串
        """
        # 由子类实现
        return str(time_value)

    def extract_variable(self, variable_name: str = None) -> Optional[xr.DataArray]:
        """
        从数据集中提取指定变量

        Args:
            variable_name: 要提取的变量名称，如果为None则从配置中获取

        Returns:
            提取的变量数据数组，如果变量不存在则返回None
        """
        if self.dataset is None:
            return None

        # 如果未指定变量名，则从配置中获取，默认为'P'
        if variable_name is None:
            try:
                from ..config import get_config
                variable_name = get_config('era5_variable', 'P')
            except ImportError:
                # 如果无法导入配置模块，使用默认值
                variable_name = 'P'

        if variable_name not in self.dataset.variables:
            return None

        return self.dataset[variable_name]

    def get_spatial_dimensions(self) -> Tuple[str, str]:
        """
        获取数据集的空间维度名称

        Returns:
            纬度和经度维度名称的元组
        """
        # 尝试常见的维度名称
        lat_candidates = ['lat', 'latitude', 'y']
        lon_candidates = ['lon', 'longitude', 'x']

        lat_dim = None
        for lat_name in lat_candidates:
            if lat_name in self.dataset.dims:
                lat_dim = lat_name
                break

        lon_dim = None
        for lon_name in lon_candidates:
            if lon_name in self.dataset.dims:
                lon_dim = lon_name
                break

        if lat_dim is None or lon_dim is None:
            raise ValueError(f"无法识别数据集的空间维度，可用维度: {list(self.dataset.dims)}")

        return lat_dim, lon_dim

    def close(self) -> None:
        """
        关闭数据集
        """
        if self.dataset is not None:
            self.dataset.close()


class DailyParser(ERA5Parser):
    """
    日降水量数据解析器
    用于解析ERA5日降水量NetCDF文件
    """

    def get_time_dimension(self) -> str:
        """
        获取时间维度名称
        
        Returns:
            时间维度名称
        """
        # 检查可能的时间维度名称
        time_candidates = ['bday', 'time', 'day']
        
        for time_name in time_candidates:
            if time_name in self.dataset.dims:
                return time_name
        
        # 如果没有找到标准时间维度，返回第一个维度
        if self.dataset.dims:
            return list(self.dataset.dims)[0]
        
        raise ValueError("无法找到时间维度")

    def format_date(self, time_value) -> str:
        """
        将bday值转换为日期字符串
        
        Args:
            time_value: bday值（从年初开始的天数）
            
        Returns:
            格式化的日期字符串 (YYYY-MM-DD)
        """
        if self.year is None:
            return f"day_{time_value}"
        
        try:
            # 将bday转换为日期
            base_date = datetime.datetime(self.year, 1, 1)
            target_date = base_date + datetime.timedelta(days=int(time_value))
            return target_date.strftime('%Y-%m-%d')
        except (ValueError, OverflowError):
            return f"{self.year}_day_{time_value}"

    def get_month_day_ranges(self) -> dict:
        """
        获取每个月份对应的天数范围（bday值）
        
        Returns:
            月份到天数范围的映射字典
        """
        if self.year is None:
            # 使用平年的天数分布
            return {
                1: (0, 30),    # 1月: 0-30天
                2: (31, 58),   # 2月: 31-58天
                3: (59, 89),   # 3月: 59-89天
                4: (90, 119),  # 4月: 90-119天
                5: (120, 150), # 5月: 120-150天
                6: (151, 180), # 6月: 151-180天
                7: (181, 211), # 7月: 181-211天
                8: (212, 242), # 8月: 212-242天
                9: (243, 272), # 9月: 243-272天
                10: (273, 303), # 10月: 273-303天
                11: (304, 333), # 11月: 304-333天
                12: (334, 365)  # 12月: 334-365天
            }
        
        # 根据具体年份计算准确的天数范围
        month_ranges = {}
        cumulative_days = 0
        
        for month in range(1, 13):
            # 计算该月的天数
            if month == 2:
                # 检查是否为闰年
                is_leap = (self.year % 4 == 0 and self.year % 100 != 0) or (self.year % 400 == 0)
                days_in_month = 29 if is_leap else 28
            elif month in [4, 6, 9, 11]:
                days_in_month = 30
            else:
                days_in_month = 31
            
            start_day = cumulative_days
            end_day = cumulative_days + days_in_month - 1
            month_ranges[month] = (start_day, end_day)
            cumulative_days += days_in_month
        
        return month_ranges


class HourlyParser(ERA5Parser):
    """
    小时降水量数据解析器
    用于解析ERA5小时降水量NetCDF文件
    """

    def get_time_dimension(self) -> str:
        """
        获取时间维度名称
        
        Returns:
            时间维度名称
        """
        # 检查可能的时间维度名称
        time_candidates = ['time', 'hour', 'valid_time']
        
        for time_name in time_candidates:
            if time_name in self.dataset.dims:
                return time_name
        
        # 如果没有找到标准时间维度，返回第一个维度
        if self.dataset.dims:
            return list(self.dataset.dims)[0]
        
        raise ValueError("无法找到时间维度")

    def format_date(self, time_value) -> str:
        """
        格式化小时数据的时间值
        
        Args:
            time_value: 时间值
            
        Returns:
            格式化的日期时间字符串
        """
        try:
            # 尝试将时间值转换为pandas时间戳
            timestamp = pd.to_datetime(time_value)
            return timestamp.strftime('%Y-%m-%d %H:%M:%S')
        except:
            return str(time_value)


def create_parser(file_path: str) -> Optional[ERA5Parser]:
    """
    根据文件路径创建合适的ERA5解析器
    
    Args:
        file_path: ERA5数据文件路径
        
    Returns:
        相应的ERA5解析器实例，如果无法创建则返回None
    """
    try:
        # 打开数据集
        dataset = xr.open_dataset(file_path)
        
        # 根据文件名判断数据类型
        filename = os.path.basename(file_path)
        
        if 'daily' in filename.lower():
            return DailyParser(dataset, file_path)
        elif 'hourly' in filename.lower():
            return HourlyParser(dataset, file_path)
        else:
            # 根据时间维度判断
            if 'bday' in dataset.dims:
                return DailyParser(dataset, file_path)
            elif 'time' in dataset.dims:
                return HourlyParser(dataset, file_path)
            else:
                # 默认使用日解析器
                return DailyParser(dataset, file_path)
                
    except Exception as e:
        print(f"创建ERA5解析器失败: {e}")
        return None